#include "linear_hmmsq.hpp"
#include "helib_atm_wrapper.hpp"
#include "log_helper.hpp"

// typedef helib::zzX PtxtType;
typedef helib::EncodedPtxt PtxtType;
typedef std::pair<long, PtxtType> diagonal_t;

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

class linear_rotate_align_ch : public constant_handle {
private:
  template <typename T> void my_init(hmmsq_collector *collector);

public:
  linear_rotate_align_ch(const helib::EncryptedArray &ea, long dim)
      : constant_handle(ea), dim_(dim) {
    bool row_is_2 = ea_.dimension() == 2 && ea_.sizeOfDimension(0) == 2;
    ASSERT_PRINTF(ea_.dimension() == 1 || row_is_2,
                  "Only support one-dimensional slots\n");
  }

  void init(hmm_metric_collector *collector) override;

  const std::vector<diagonal_t> &ra_left() const { return rotate_align_left_; }
  const std::vector<diagonal_t> &ra_up() const { return rotate_align_up_; }

private:
  long dim_;

  std::vector<diagonal_t> rotate_align_left_;
  std::vector<diagonal_t> rotate_align_up_;
};

class linear_shift_compute_ch : public constant_handle {
private:
  template <typename T> void my_init(hmmsq_collector *collector);

public:
  linear_shift_compute_ch(const helib::EncryptedArray &ea, long dim)
      : constant_handle(ea), dim_(dim) {
    bool row_is_2 = ea_.dimension() == 2 && ea_.sizeOfDimension(0) == 2;
    ASSERT_PRINTF(ea_.dimension() == 1 || row_is_2,
                  "Only support one-dimensional slots\n");
  }

  void init(hmm_metric_collector *collector) override;

  const std::vector<diagonal_t> &shift_left(std::size_t i) const {
    return shift_left_[i];
  }
  const std::vector<diagonal_t> &shift_up(std::size_t i) const {
    return shift_up_[i];
  }

private:
  long dim_;

  std::vector<std::vector<diagonal_t>> shift_left_;
  std::vector<std::vector<diagonal_t>> shift_up_;
};

class lhmmsq_constant_multipler : public hmm_constant_multipler {
public:
  lhmmsq_constant_multipler(const helib::EncryptedArray &ea,
                            hmm_metric_collector *collector, long dim)
      : dim_(dim), racm_(new linear_rotate_align_ch(ea, dim)),
        sccm_(new linear_shift_compute_ch(ea, dim)) {
    hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
    ASSERT_PRINTF(hmmsq != nullptr,
                  "Failed to static_cast to `hmmsq_collector *`\n");

    hmmsq->start_encoding_ra(); //! Encoding in RotateAlign - Start
    racm_->init(collector);
    hmmsq->stop_encoding_ra(); //! Encoding in RotateAlign - Stop

    hmmsq->start_encoding_sc(); //! Encoding in ShiftCompute - Start
    sccm_->init(collector);
    hmmsq->stop_encoding_sc(); //! Encoding in ShiftCompute - Stop
  }
  ~lhmmsq_constant_multipler() = default;

  const std::vector<diagonal_t> &ra_left() const { return racm_->ra_left(); }
  const std::vector<diagonal_t> &ra_up() const { return racm_->ra_up(); }

  const std::vector<diagonal_t> &shift_left(std::size_t i) const {
    return sccm_->shift_left(i);
  }
  const std::vector<diagonal_t> &shift_up(std::size_t i) const {
    return sccm_->shift_up(i);
  }

  void set_no_encoding_reuse(bool val) override { no_encoding_reuse_ = val; }

  long dim() const { return dim_; }

private:
  // FIXME:
  long dim_;

  linear_rotate_align_ch *racm_ = nullptr;
  linear_shift_compute_ch *sccm_ = nullptr;
};

hmm_constant_multipler *
build_lhmmsq_constant_multipler(const helib::EncryptedArray &ea,
                                hmm_metric_collector *collector, long dim) {
  return new lhmmsq_constant_multipler(ea, collector, dim);
}

template <typename T>
static inline void set_diagonal(std::vector<T> &diag, long idx, long dim,
                                long inner_offset = -1,
                                long outter_offset = 0) {
  ASSERT_PRINTF(inner_offset < dim && outter_offset < dim, "Invalid offset\n");

  bool is_zero = true;
  const long dim_sq = dim * dim;
  for (long i = 0; i < diag.size(); i++) {
    long j = (i + idx) % diag.size();
    if (i >= dim_sq || j >= dim_sq) {
      continue;
    }
    // Divide dim_sq * dim_sq by dim * dim
    long blk_i = i / dim;
    long blk_j = j / dim;
    if ((blk_j - blk_i + dim) % dim != outter_offset) {
      continue;
    }
    long ii = i % dim;
    long jj = j % dim;
    if ((jj - ii + dim) % dim == (inner_offset == -1 ? blk_i : inner_offset)) {
      diag[i] = static_cast<T>(1);
      is_zero = true;
    }
  }

  ASSERT_PRINTF(is_zero != false, "diagonal %ld is zero\n", idx);
}

template <typename T>
void linear_rotate_align_ch::my_init(hmmsq_collector *hmmsq) {
  // left
  const long dim_sq = dim_ * dim_;
  const long nslot = ea_.size();
  const long bound = dim_ - 1;
  rotate_align_left_.resize(2 * bound + 1);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (long i = 0; i < 2 * bound + 1; i++) {
    auto tid = omp_get_thread_num();

    //! [-bound, bound]: 2*dim_-1
    //! nslot >= dim_*dim_
    //! dim_*dim_ >= 2*dim_-1 <=> (dim_-1)^2
    //! nslot >= 2*dim_-1
    //! idx cannot be the same
    long k = -bound + i;
    const long idx = (k + nslot) % nslot;

    hmmsq->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
    std::vector<T> diag(nslot, 0);
    set_diagonal(diag, idx, dim_);
    hmmsq->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Stop

    hmmsq->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
    PtxtType poly;
    ea_.encode(poly, diag);
    hmmsq->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Stop
    rotate_align_left_[i] = {idx, poly};
  }

  // up
  //* (i,j) -> (ii, jj)
  //* i*D+j -> ii*D+jj, where
  //* ii = (i-j+D)%D, jj = j
  hmmsq->start_encoding_ra_raw2vec(); //! Raw 2 Vec - Start

  // key: diagnoal_idx; value: diagnoal
  std::unordered_map<long, std::vector<T>> tmp;
  for (int k = 0; k < dim_sq; k++) {
    long i = k / dim_;
    long j = k % dim_;
    long kk = ((i - j + dim_) % dim_) * dim_ + j;
    // (k, k) -> (kk, k)
    const long idx = (k - kk + nslot) % nslot;
    auto res = tmp.find(idx);
    if (res != tmp.end()) {
      res->second[kk] = 1;
    } else {
      tmp.insert({idx, std::vector<T>(nslot, static_cast<T>(0))});
      tmp[idx][kk] = 1;
    }
  }
  hmmsq->stop_encoding_ra_raw2vec(); //! Raw 2 Vec - Stop

  hmmsq->start_encoding_ra_vec2ptxt(); //! Vec 2 Ptxt - Start
  int idx = -1;
  rotate_align_up_.resize(tmp.size());
#pragma omp parallel
#pragma omp single
  for (auto &p : tmp) {
    idx++;
#pragma omp task
    {
      PtxtType poly;
      ea_.encode(poly, p.second);
      rotate_align_up_[idx] = {p.first, poly};
    }
  }
  hmmsq->stop_encoding_ra_vec2ptxt(); //! Vec 2 Ptxt - Stop
}

void linear_rotate_align_ch::init(hmm_metric_collector *collector) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  if (ea_.isCKKS()) {
    my_init<double>(hmmsq);
  } else {
    my_init<long>(hmmsq);
  }
}

template <typename T>
void linear_shift_compute_ch::my_init(hmmsq_collector *hmmsq) {
  const long nslot = ea_.size();
  //! shift_left_[0] && shift_left_[0] are empty
  shift_left_.resize(dim_);
  shift_up_.resize(dim_);

#pragma omp parallel
#pragma omp single
  {
    // shift left
#pragma omp taskloop nogroup
    for (long i = 1; i < dim_; i++) {
      auto tid = omp_get_thread_num();

      //! 1-dim_+nslot !=  1+nslot (mod nslot) when dim_>1
      hmmsq->start_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Start
      const long left1_idx = (i - dim_ + nslot) % nslot;
      std::vector<T> left1(nslot, 0);
      set_diagonal(left1, left1_idx, dim_, i);
      hmmsq->stop_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Stop

      hmmsq->start_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      PtxtType poly;
      ea_.encode(poly, left1);
      shift_left_[i].emplace_back(left1_idx, poly);
      hmmsq->stop_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Stop

      const long left2_idx = (i + nslot) % nslot;
      if (left1_idx != left2_idx) {
        hmmsq->start_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Start
        std::vector<T> left2(nslot, 0);
        set_diagonal(left2, left2_idx, dim_, i);
        hmmsq->stop_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Stop

        hmmsq->start_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Start
        PtxtType poly;
        ea_.encode(poly, left2);
        shift_left_[i].emplace_back(left2_idx, poly);
        hmmsq->stop_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Stop
      }
    }

    // shift up
#pragma omp taskloop nogroup
    for (long i = 1; i < dim_; i++) {
      auto tid = omp_get_thread_num();

      hmmsq->start_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Start
      const long up_idx = i * dim_;
      std::vector<T> up(nslot, 0);
      set_diagonal(up, up_idx, dim_, 0, i);
      hmmsq->stop_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Stop

      hmmsq->start_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      PtxtType poly;
      ea_.encode(poly, up);
      shift_up_[i].emplace_back(up_idx, poly);
      hmmsq->stop_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Stop
      if (dim_ * dim_ != nslot) {
        hmmsq->start_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Start
        const long up_idx = i * dim_ + nslot - dim_ * dim_;
        std::vector<T> up(nslot, 0);
        set_diagonal(up, up_idx, dim_, 0, i);
        hmmsq->stop_encoding_sc_raw2vec(tid); //! Raw 2 Vec - Stop

        hmmsq->start_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Start
        PtxtType poly;
        ea_.encode(poly, up);
        shift_up_[i].emplace_back(up_idx, poly);
        hmmsq->stop_encoding_sc_vec2ptxt(tid); //! Vec 2 Ptxt - Stop
      }
    }
  }
}

void linear_shift_compute_ch::init(hmm_metric_collector *collector) {

  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  if (ea_.isCKKS()) {
    my_init<double>(hmmsq);
  } else {
    my_init<long>(hmmsq);
  }
}

void lhmm_sq(hmm_metric_collector *collector, hmm_constant_multipler *cm,
             const helib::EncryptedArray &ea, helib::Ctxt &AB,
             const helib::Ctxt &A, const helib::Ctxt &B) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<lhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `lhmmsq_constant_multipler *`\n");

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));
  hmmsq->init_err = std::max(A.errorBound(), B.errorBound());

  bool row_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(0) == 2;
  bool col_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(1) == 2;
  long dim_idx = row_is_2 ? 1 : 0;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);
  // std::vector<helib::Ctxt> AB_partials(hmmsqcm->dim(), helib::Ctxt(pk));

  //! RotateAlign - Start
  hmmsq->start_computing_ra();
  hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(A, dim_idx, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_B(B, dim_idx, ea);
  hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

  // RotateAlign for A
  for (auto &p : hmmsqcm->ra_left()) {
    const auto offset = p.first;
    const auto constant = p.second;

    hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
    auto A_i = hoisted_A.automorph(-1l * offset);
    hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmmsq->start_computing_ra_scale(); //! RA.Scale - Start
    A_i->multByConstant(constant);
    hmmsq->stop_computing_ra_scale(); //! RA.Scale - Stop

    hmmsq->start_computing_ra_add(); //! RA.Add - Start
    Ap += *A_i;
    hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
  }

  // RotateAlign for B
  for (auto &p : hmmsqcm->ra_up()) {
    const auto offset = p.first;
    const auto constant = p.second;

    hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
    auto B_i = hoisted_B.automorph(-1l * offset);
    hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmmsq->start_computing_ra_scale(); //! RA.Scale - Start
    B_i->multByConstant(constant);
    hmmsq->stop_computing_ra_scale(); //! RA.Scale - Stop

    hmmsq->start_computing_ra_add(); //! RA.Add - Start
    Bp += *B_i;
    hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
  }
  hmmsq->stop_computing_ra();
  hmmsq->ra_noise =
      std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) - hmmsq->init_noise;
  hmmsq->ra_err = std::max(Ap.errorBound(), Bp.errorBound());

  //! ShiftCompute - Start
  hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));

  hmmsq->start_computing_sc();
  hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
  helib::GeneralAutomorphPrecon_FULL hoisted_Ap(Ap, dim_idx, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_Bp(Bp, dim_idx, ea);
  hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop

  hmmsq->start_computing_sc_mult(); //! SC.Mult - Start
  // AB_partials[0] = Ap;
  // AB_partials[0] *= Bp;
  AB_temp = Ap;
  AB_temp *= Bp;
  // std::cout << "0: " << AB_temp.errorBound() << std::endl;
  hmmsq->stop_computing_sc_mult(); //! SC.Mult - Stop

  for (long i = 1; i < hmmsqcm->dim(); i++) {
    std::shared_ptr<helib::Ctxt> A_i;
    std::shared_ptr<helib::Ctxt> B_i;
    // Shift A
    const auto &shift_left_const = hmmsqcm->shift_left(i);
    for (auto it = shift_left_const.begin(); it != shift_left_const.end();
         it++) {
      const auto offset = it->first;
      const auto constant = it->second;

      hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
      auto A_ii = hoisted_Ap.automorph(-1l * offset);
      hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop

      hmmsq->start_computing_sc_scale(); //! SC.Scale - Start
      A_ii->multByConstant(constant);
      hmmsq->stop_computing_sc_scale(); //! SC.Scale - Stop

      if (it == shift_left_const.begin()) {
        A_i = A_ii;
      } else {
        hmmsq->start_computing_sc_add(); //! SC.Add - Start
        *A_i += *A_ii;
        hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
      }
    }
    // Shift B
    const auto &shift_up_const = hmmsqcm->shift_up(i);
    for (auto it = shift_up_const.begin(); it != shift_up_const.end(); it++) {
      const auto offset = it->first;
      const auto constant = it->second;

      hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
      auto B_ii = hoisted_Bp.automorph(-1l * offset);
      hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop

      hmmsq->start_computing_sc_scale(); //! SC.Scale - Start
      B_ii->multByConstant(constant);
      hmmsq->stop_computing_sc_scale(); //! SC.Scale - Stop

      if (it == shift_up_const.begin()) {
        B_i = B_ii;
      } else {
        hmmsq->start_computing_sc_add(); //! SC.Add - Start
        *B_i += *B_ii;
        hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
      }
    }
    hmmsq->sc_rot_err = std::max(A_i->errorBound(), B_i->errorBound());

    hmmsq->start_computing_sc_mult(); //! SC.Mult - Start
    A_i->multiplyBy(*B_i);
    hmmsq->stop_computing_sc_mult(); //! SC.Mult - Stop
    hmmsq->sc_hmul_err = std::max(A_i->errorBound(), B_i->errorBound());

    // AB_partials[i] = std::move(*A_i);

    // hmmsq->start_computing_sc_add(); //! SC.Add - Start
    AB_temp += *A_i;
    hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
    // std::cout << i << ": " << AB_temp.errorBound() << std::endl;
  }
  hmmsq->stop_computing_sc();
  //! ShiftCompute - Completed

  // hmmsq->start_computing_sc_add(); //! SC.Add - Start
  // for (int stride = 1; 2 * stride <= AB_partials.size(); stride *= 2) {
  //   for (int i = 0; 2 * stride * i + stride < AB_partials.size(); i++) {
  //     AB_partials[2 * stride * i] += AB_partials[2 * stride * i + stride];
  //     std::cout << stride << ":" << i << ": "
  //               << AB_partials[2 * stride * i].errorBound() << std::endl;
  //   }
  // }
  // hmmsq->stop_computing_sc_add(); //! SC.Add - Stop

  // AB = std::move(AB_partials[0]);
  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
  hmmsq->final_err = AB.errorBound();
}

void lhmm_sq_mt(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                int threads, const helib::EncryptedArray &ea, helib::Ctxt &AB,
                const helib::Ctxt &A, const helib::Ctxt &B) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<lhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `lhmmsq_constant_multipler *`\n");

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));
  hmmsq->init_err = std::max(A.errorBound(), B.errorBound());

  bool row_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(0) == 2;
  bool col_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(1) == 2;
  long dim_idx = row_is_2 ? 1 : 0;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;

  //! RotateAlign - Start
  const auto &ra_left_const = hmmsqcm->ra_left();
  const auto &ra_up_const = hmmsqcm->ra_up();

#pragma omp parallel
#pragma omp single
  {
    hmmsq->start_computing_ra();
#pragma omp task
    hoisted_A =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, dim_idx, ea);
#pragma omp task
    hoisted_B =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, dim_idx, ea);
#pragma omp taskwait //! Explicit synchronization

    // RotateAlign for A
#pragma omp taskloop nogroup
    for (const auto &p : hmmsqcm->ra_left()) {
      const auto offset = p.first;
      const auto constant = p.second;

      auto A_i = hoisted_A->automorph(-1l * offset);
      A_i->multByConstant(constant);

#pragma omp critical
      Ap += *A_i;
    }

    // RotateAlign for B
#pragma omp taskloop nogroup
    for (const auto &p : hmmsqcm->ra_up()) {
      const auto offset = p.first;
      const auto constant = p.second;

      auto B_i = hoisted_B->automorph(-1l * offset);
      B_i->multByConstant(constant);

#pragma omp critical
      Bp += *B_i;
    }
#pragma omp taskwait //! Explicit synchronization
    hmmsq->stop_computing_ra();
    hmmsq->ra_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
                      hmmsq->init_noise;
    hmmsq->ra_err = std::max(Ap.errorBound(), Bp.errorBound());
    //! RotateAlign - Completed

    //! ShiftCompute - Start
    hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
    hmmsq->start_computing_sc();
#pragma omp task
    hoisted_Ap =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, dim_idx, ea);
#pragma omp task
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, dim_idx, ea);
#pragma omp taskwait //! Explicit synchronization

#pragma omp task
    {
      AB_temp = Ap;
      AB_temp *= Bp;
    }

#pragma omp taskloop
    for (long i = 1; i < hmmsqcm->dim(); i++) {
      std::shared_ptr<helib::Ctxt> A_i;
      std::shared_ptr<helib::Ctxt> B_i;
      // Shift A
      const auto &shift_left_const = hmmsqcm->shift_left(i);
      for (auto it = shift_left_const.begin(); it != shift_left_const.end();
           it++) {
        const auto offset = it->first;
        const auto constant = it->second;

        auto A_ii = hoisted_Ap->automorph(-1l * offset);
        A_ii->multByConstant(constant);

        if (it == shift_left_const.begin()) {
          A_i = A_ii;
        } else {
          *A_i += *A_ii;
        }
      }
      // Shift B
      const auto &shift_up_const = hmmsqcm->shift_up(i);
      for (auto it = shift_up_const.begin(); it != shift_up_const.end(); it++) {
        const auto offset = it->first;
        const auto constant = it->second;

        auto B_ii = hoisted_Bp->automorph(-1l * offset);
        B_ii->multByConstant(constant);

        if (it == shift_up_const.begin()) {
          B_i = B_ii;
        } else {
          *B_i += *B_ii;
        }
      }
      if (i == 1) {
        hmmsq->sc_rot_err = std::max(A_i->errorBound(), B_i->errorBound());
      }
      A_i->multiplyBy(*B_i);
      if (i == 1) {
        hmmsq->sc_hmul_err = std::max(A_i->errorBound(), B_i->errorBound());
      }
#pragma omp critical
      AB_temp += *A_i;
    } //! omp taskloop: Implicit synchronization
    hmmsq->stop_computing_sc();
    //! ShiftCompute - Completed
  }

  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
  hmmsq->final_err = AB.errorBound();
}